from Python import retinaface_native
import numpy as np
import ctypes
import cv2


class RetinaFaceDetection():
    def __init__(self, model_dir):
        model_dir_path = retinaface_native.String(bytes(model_dir, encoding="utf8"))
        self.session = retinaface_native.InitDetectorSession(model_dir_path)

    def detection(self, image, level=320):
        mem = ctypes.POINTER(ctypes.c_float)()
        size = ctypes.c_int(0)
        h, w, _ = image.shape
        retinaface_native.FaceDetection(self.session, image.tostring(), w, h, level, ctypes.byref(mem),
                                        ctypes.byref(size))
        res = np.ctypeslib.as_array((ctypes.c_float * 15 * size.value).from_address(ctypes.addressof(mem.contents)))
        numpyed_bak = res.reshape(size.value, 15).copy()
        retinaface_native.FreeFloat(ctypes.byref(mem))
        return numpyed_bak

# det = RetinaFaceDetection('/Users/yh-mac/my_project/bishe/retinaface_cpp_implentation/models')
# # img = cv2.imread("/Users/yh-mac/Desktop/mz.jpg")
#
# cap = cv2.VideoCapture(0)
# # 循环读取每一帧
# while (cap.isOpened()):
#     ret, frame = cap.read()
#     res = np.asarray(det.detection(frame, 320), dtype=int)
#     for item in res:
#         bbox = item[1:5]
#         x1, y1, x2, y2 = bbox
#         point = item[5:]
#         cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
#         for x, y in point.reshape(5, 2):
#             cv2.line(frame, (x, y), (x, y), (200, 100, 80), 5)
#     cv2.imshow("frame", frame)  # 显示图片
#     if cv2.waitKey(1) & 0xFF == ord("q"):  # 按q退出
#         break
# # print(res)
